{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "initial_id",
   "metadata": {
    "collapsed": true,
    "ExecuteTime": {
     "end_time": "2024-03-04T13:59:52.983596Z",
     "start_time": "2024-03-04T13:59:49.250397Z"
    }
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import torch\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "stock = 'hs300'"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-04T13:59:52.999599Z",
     "start_time": "2024-03-04T13:59:52.986596Z"
    }
   },
   "id": "7eb19cde5f1714e7",
   "execution_count": 2
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "basic_feature = [\"open\", \"close\", \"high\", \"low\", \"volume\", \"turnover\"]\n",
    "# basic_feature = [\"open\", \"close\", \"high\", \"low\", \"volume\"]"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-04T13:59:53.015603Z",
     "start_time": "2024-03-04T13:59:53.001600Z"
    }
   },
   "id": "c55ce71f087a267a",
   "execution_count": 3
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "target_return_span = 1\n",
    "target = f\"return+{target_return_span}\""
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-04T13:59:53.031607Z",
     "start_time": "2024-03-04T13:59:53.018604Z"
    }
   },
   "id": "b9ed4d57248ee784",
   "execution_count": 4
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "time_span = 60\n",
    "# stock_date = ['2018-01-01', '2021-12-31', '2022-01-02', '2022-10-04',  '2022-10-05', '2023-12-31']\n",
    "stock_date = ['2018-01-01', '2021-12-31', '2022-01-02', '2022-09-29',  '2022-09-30', '2023-12-31']"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-04T13:59:53.047612Z",
     "start_time": "2024-03-04T13:59:53.034607Z"
    }
   },
   "id": "e6d6802c16bec108",
   "execution_count": 5
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "tic_df_list = []\n",
    "for tic_path in os.listdir(f\"./{stock}\"):\n",
    "    # tic_df = pd.read_csv(f\"./{stock}/{tic_path}\")[[\"date\", \"open\", \"close\", \"high\", \"low\", \"volume\", \"tic\"]]\n",
    "    tic_df = pd.read_csv(f\"./{stock}/{tic_path}\")[[\"date\", \"open\", \"close\", \"high\", \"low\", \"volume\", \"turnover\", \"tic\"]]\n",
    "    tic_df[target] = tic_df.close.pct_change(target_return_span).shift(-1 * target_return_span)\n",
    "    tic_df_list.append(tic_df)\n",
    "\n",
    "tic_target_df = pd.concat(tic_df_list).sort_values(by=\"date\").dropna()\n",
    "tic_target_df.date = pd.DatetimeIndex(tic_target_df.date)"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-04T13:59:53.889349Z",
     "start_time": "2024-03-04T13:59:53.049612Z"
    }
   },
   "id": "524612c4c9b16c3e",
   "execution_count": 6
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "data": {
      "text/plain": "           date   open  close   high    low       volume      turnover  \\\n0    2018-01-02  13.35  13.70  13.93  13.32  208159255.0  2.856544e+09   \n0    2018-01-02   4.65   4.72   4.84   4.61  257111210.0  1.216979e+09   \n0    2018-01-02   7.13   7.18   7.23   7.12   49500856.0  3.559108e+08   \n0    2018-01-02  17.05  17.24  17.29  17.01   97602162.0  1.678255e+09   \n0    2018-01-02  29.55  29.02  30.66  28.92   36416515.0  1.088832e+09   \n...         ...    ...    ...    ...    ...          ...           ...   \n1455 2023-12-28  14.52  14.74  14.75  14.48   49900548.0  7.321915e+08   \n1455 2023-12-28  19.47  20.20  20.30  19.33   57302141.0  1.147868e+09   \n1455 2023-12-28   7.96   7.82   8.05   7.74   71736980.0  5.615205e+08   \n1455 2023-12-28   9.11   9.45   9.47   9.08  166159184.0  1.550257e+09   \n1455 2023-12-28  20.58  21.33  21.46  20.41   16601045.0  3.509709e+08   \n\n            tic  return+1  \n0     000001.SZ -0.027007  \n0     000425.SZ -0.006356  \n0     601169.SH  0.011142  \n0     601166.SH -0.002320  \n0     601155.SH -0.000689  \n...         ...       ...  \n1455  601166.SH  0.099729  \n1455  002142.SZ -0.004455  \n1455  600011.SH -0.015345  \n1455  000001.SZ -0.006349  \n1455  601877.SH  0.008439  \n\n[203840 rows x 9 columns]",
      "text/html": "<div>\n<style scoped>\n    .dataframe tbody tr th:only-of-type {\n        vertical-align: middle;\n    }\n\n    .dataframe tbody tr th {\n        vertical-align: top;\n    }\n\n    .dataframe thead th {\n        text-align: right;\n    }\n</style>\n<table border=\"1\" class=\"dataframe\">\n  <thead>\n    <tr style=\"text-align: right;\">\n      <th></th>\n      <th>date</th>\n      <th>open</th>\n      <th>close</th>\n      <th>high</th>\n      <th>low</th>\n      <th>volume</th>\n      <th>turnover</th>\n      <th>tic</th>\n      <th>return+1</th>\n    </tr>\n  </thead>\n  <tbody>\n    <tr>\n      <th>0</th>\n      <td>2018-01-02</td>\n      <td>13.35</td>\n      <td>13.70</td>\n      <td>13.93</td>\n      <td>13.32</td>\n      <td>208159255.0</td>\n      <td>2.856544e+09</td>\n      <td>000001.SZ</td>\n      <td>-0.027007</td>\n    </tr>\n    <tr>\n      <th>0</th>\n      <td>2018-01-02</td>\n      <td>4.65</td>\n      <td>4.72</td>\n      <td>4.84</td>\n      <td>4.61</td>\n      <td>257111210.0</td>\n      <td>1.216979e+09</td>\n      <td>000425.SZ</td>\n      <td>-0.006356</td>\n    </tr>\n    <tr>\n      <th>0</th>\n      <td>2018-01-02</td>\n      <td>7.13</td>\n      <td>7.18</td>\n      <td>7.23</td>\n      <td>7.12</td>\n      <td>49500856.0</td>\n      <td>3.559108e+08</td>\n      <td>601169.SH</td>\n      <td>0.011142</td>\n    </tr>\n    <tr>\n      <th>0</th>\n      <td>2018-01-02</td>\n      <td>17.05</td>\n      <td>17.24</td>\n      <td>17.29</td>\n      <td>17.01</td>\n      <td>97602162.0</td>\n      <td>1.678255e+09</td>\n      <td>601166.SH</td>\n      <td>-0.002320</td>\n    </tr>\n    <tr>\n      <th>0</th>\n      <td>2018-01-02</td>\n      <td>29.55</td>\n      <td>29.02</td>\n      <td>30.66</td>\n      <td>28.92</td>\n      <td>36416515.0</td>\n      <td>1.088832e+09</td>\n      <td>601155.SH</td>\n      <td>-0.000689</td>\n    </tr>\n    <tr>\n      <th>...</th>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n    </tr>\n    <tr>\n      <th>1455</th>\n      <td>2023-12-28</td>\n      <td>14.52</td>\n      <td>14.74</td>\n      <td>14.75</td>\n      <td>14.48</td>\n      <td>49900548.0</td>\n      <td>7.321915e+08</td>\n      <td>601166.SH</td>\n      <td>0.099729</td>\n    </tr>\n    <tr>\n      <th>1455</th>\n      <td>2023-12-28</td>\n      <td>19.47</td>\n      <td>20.20</td>\n      <td>20.30</td>\n      <td>19.33</td>\n      <td>57302141.0</td>\n      <td>1.147868e+09</td>\n      <td>002142.SZ</td>\n      <td>-0.004455</td>\n    </tr>\n    <tr>\n      <th>1455</th>\n      <td>2023-12-28</td>\n      <td>7.96</td>\n      <td>7.82</td>\n      <td>8.05</td>\n      <td>7.74</td>\n      <td>71736980.0</td>\n      <td>5.615205e+08</td>\n      <td>600011.SH</td>\n      <td>-0.015345</td>\n    </tr>\n    <tr>\n      <th>1455</th>\n      <td>2023-12-28</td>\n      <td>9.11</td>\n      <td>9.45</td>\n      <td>9.47</td>\n      <td>9.08</td>\n      <td>166159184.0</td>\n      <td>1.550257e+09</td>\n      <td>000001.SZ</td>\n      <td>-0.006349</td>\n    </tr>\n    <tr>\n      <th>1455</th>\n      <td>2023-12-28</td>\n      <td>20.58</td>\n      <td>21.33</td>\n      <td>21.46</td>\n      <td>20.41</td>\n      <td>16601045.0</td>\n      <td>3.509709e+08</td>\n      <td>601877.SH</td>\n      <td>0.008439</td>\n    </tr>\n  </tbody>\n</table>\n<p>203840 rows × 9 columns</p>\n</div>"
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset_df = tic_target_df\n",
    "dataset_df"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-04T13:59:53.921356Z",
     "start_time": "2024-03-04T13:59:53.890349Z"
    }
   },
   "id": "6eb5193ba4cdb196",
   "execution_count": 7
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'pandas.core.frame.DataFrame'>\n",
      "Index: 203840 entries, 0 to 1455\n",
      "Data columns (total 9 columns):\n",
      " #   Column    Non-Null Count   Dtype         \n",
      "---  ------    --------------   -----         \n",
      " 0   date      203840 non-null  datetime64[ns]\n",
      " 1   open      203840 non-null  float64       \n",
      " 2   close     203840 non-null  float64       \n",
      " 3   high      203840 non-null  float64       \n",
      " 4   low       203840 non-null  float64       \n",
      " 5   volume    203840 non-null  float64       \n",
      " 6   turnover  203840 non-null  float64       \n",
      " 7   tic       203840 non-null  object        \n",
      " 8   return+1  203840 non-null  float64       \n",
      "dtypes: datetime64[ns](1), float64(7), object(1)\n",
      "memory usage: 15.6+ MB\n"
     ]
    }
   ],
   "source": [
    "dataset_df.info()"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-04T13:59:53.953365Z",
     "start_time": "2024-03-04T13:59:53.925357Z"
    }
   },
   "id": "e0e89d23ca47effd",
   "execution_count": 8
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "data": {
      "text/plain": "1456"
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset_df['date'].nunique()"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-04T13:59:53.984876Z",
     "start_time": "2024-03-04T13:59:53.955869Z"
    }
   },
   "id": "d28c862f7a53e9e4",
   "execution_count": 9
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "train_df = dataset_df[(dataset_df.date >= stock_date[0]) & (dataset_df.date <= stock_date[1])]\n",
    "val_df = dataset_df[(dataset_df.date >= stock_date[2]) & (dataset_df.date <= stock_date[3])]\n",
    "test_df = dataset_df[(dataset_df.date >= stock_date[4]) & (dataset_df.date <= stock_date[5])]"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-04T13:59:54.032887Z",
     "start_time": "2024-03-04T13:59:53.988877Z"
    }
   },
   "id": "3a11832ed7bccac1",
   "execution_count": 10
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "302   973   181\n"
     ]
    }
   ],
   "source": [
    "print(test_df['date'].nunique(), ' ', train_df['date'].nunique(), ' ', val_df['date'].nunique())"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-04T13:59:54.048892Z",
     "start_time": "2024-03-04T13:59:54.034888Z"
    }
   },
   "id": "c9ecf684158c0914",
   "execution_count": 11
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "tic_list = dataset_df['tic'].unique()"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-04T13:59:54.080403Z",
     "start_time": "2024-03-04T13:59:54.052893Z"
    }
   },
   "id": "3033c23fbd996c84",
   "execution_count": 12
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "def df_2_array(dataset_df, feat_col, target, type):\n",
    "    dataset_feat = []\n",
    "    dataset_ret = []\n",
    "    dataset_price = []\n",
    "    for tic in tic_list:\n",
    "        df = dataset_df[dataset_df.tic == tic]\n",
    "        feat = df[feat_col].to_numpy()\n",
    "        ret = df[target].to_numpy()\n",
    "        price = df['close'].to_numpy()\n",
    "        stock_feat = []\n",
    "        stock_ret = []\n",
    "        stock_price = []\n",
    "        for i in range(time_span, feat.shape[0]):\n",
    "            stock_feat.append(feat[i-time_span : i])\n",
    "            stock_ret.append(ret[i])\n",
    "            stock_price.append(price[i])\n",
    "        stock_feat = np.array(stock_feat)\n",
    "        stock_ret = np.array(stock_ret)\n",
    "        stock_price = np.array(stock_price)\n",
    "    \n",
    "        dataset_feat.append(stock_feat)\n",
    "        dataset_ret.append(stock_ret)\n",
    "        dataset_price.append(stock_price)\n",
    "\n",
    "    dataset_feat = np.array(dataset_feat).transpose((1, 2, 0, 3))\n",
    "    dataset_ret = np.array(dataset_ret).transpose((1, 0))\n",
    "    dataset_price = np.array(dataset_price).transpose((1, 0))\n",
    "    \n",
    "    dataset_feat_tensor = torch.tensor(dataset_feat, dtype=torch.float)\n",
    "    dataset_ret_tensor = torch.tensor(dataset_ret, dtype=torch.float)\n",
    "    dataset_price_tensor = torch.tensor(dataset_price, dtype=torch.float)\n",
    "    \n",
    "    torch.save(dataset_feat_tensor, f\"./dataset/{stock}/{type}/feat.pt\")\n",
    "    torch.save(dataset_ret_tensor, f\"./dataset/{stock}/{type}/ret.pt\")\n",
    "    torch.save(dataset_price_tensor, f\"./dataset/{stock}/{type}/price.pt\")    \n",
    "    \n",
    "    return dataset_feat, dataset_ret, dataset_price"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-04T13:59:54.096407Z",
     "start_time": "2024-03-04T13:59:54.082404Z"
    }
   },
   "id": "e74a978832316fe8",
   "execution_count": 13
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "_ = df_2_array(train_df, basic_feature, target, \"train\")\n",
    "_ = df_2_array(val_df, basic_feature, target, \"val\")\n",
    "dataset_feat, dataset_ret, dataset_price = df_2_array(test_df, basic_feature, target, \"test\")"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-04T13:59:57.209121Z",
     "start_time": "2024-03-04T13:59:54.098408Z"
    }
   },
   "id": "e8cfee4380b2f6c5",
   "execution_count": 14
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "data": {
      "text/plain": "(242, 60, 140, 6)"
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset_feat.shape"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-04T13:59:57.225124Z",
     "start_time": "2024-03-04T13:59:57.211121Z"
    }
   },
   "id": "971454b4bb5d094",
   "execution_count": 15
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "data": {
      "text/plain": "(242, 140)"
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset_ret.shape"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-04T13:59:57.241128Z",
     "start_time": "2024-03-04T13:59:57.227125Z"
    }
   },
   "id": "bec04048b374e1f4",
   "execution_count": 16
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
